import base64
import json
import time

import argparse
import gradio as gr
import os 

import sys

sys.path.insert(0, '/mnt/sfs/asr/code/osum_xlgeng/examples/wenetspeech/tts')
from gxl_ai_utils.utils import utils_file
try:
    sys.path.insert(1, '/mnt/sfs/asr/code/osum_xlgeng/')
    from wenet.utils.init_tokenizer import init_tokenizer
    from gxl_ai_utils.config.gxl_config import GxlNode
    from wenet.utils.init_model import init_model
    import logging
    import librosa
    import torch
    import torchaudio
except ImportError:
    print("导入error!")
    pass
is_npu = True

try:
    import torch_npu
except ImportError:
    is_npu = False
    print("torch_npu is not available. if you want to use npu, please install it.")


# gpu_id = 0

def get_args():
    parser = argparse.ArgumentParser(description='recognize with your model')
    parser.add_argument('--config', required=True, help='config file')
    parser.add_argument('--test_data', required=True, help='test data file')
    parser.add_argument('--gpu_id',
                        type=int,
                        default=-1,
                        help='gpu id for this rank, -1 for cpu')
    parser.add_argument('--checkpoint', required=True, help='checkpoint model')
    parser.add_argument('--result_dir', required=True, help='asr result file')

    args = parser.parse_args()
    print(args)
    return args

def main():
    args = get_args()
    logging.basicConfig(level=logging.DEBUG,
                        format='%(asctime)s %(levelname)s %(message)s')
    configs = utils_file.load_dict_from_yaml(args.config)
    model, configs = init_model(args, configs)
    if is_npu:
        device = torch.device(f'npu:{args.gpu_id}')
    else:
        device =torch.device(f'cuda:{args.gpu_id}')
    model = model.to(device)
    tokenizer = init_tokenizer(configs)
    print(model)
    with open(args.test_data,"r") as f:
        f = f.readlines()
        for line in f:
            key,input_wav_path = line.strip().split()
            output_wav_path=os.path.join(args.result_dir, f"output_{key}.wav")
            timestamp_ms = int(time.time() * 1000)
            now_file_tmp_path_resample = f'~/.cache/.temp/{timestamp_ms}_resample.wav'
            do_resample(input_wav_path, now_file_tmp_path_resample)
            input_wav_path = now_file_tmp_path_resample
            waveform, sample_rate = torchaudio.load(input_wav_path)
            waveform = waveform.squeeze(0)  # (channel=1, sample) -> (sample,)
            window = torch.hann_window(400)
            stft = torch.stft(waveform,
                            400,
                            160,
                            window=window,
                            return_complex=True)
            magnitudes = stft[..., :-1].abs() ** 2

            filters = torch.from_numpy(
                librosa.filters.mel(sr=sample_rate,
                                    n_fft=400,
                                    n_mels=80))
            mel_spec = filters @ magnitudes

            # NOTE(xcsong): https://github.com/openai/whisper/discussions/269
            log_spec = torch.clamp(mel_spec, min=1e-10).log10()
            log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
            log_spec = (log_spec + 4.0) / 4.0
            feat = log_spec.transpose(0, 1)
            feat_lens = torch.tensor([feat.shape[0]], dtype=torch.int64).to(device)
            feat = feat.unsqueeze(0).to(device)
            prompt = "先根据语音输入，直接以文字形式进行回答或对话，接着再生成语音token。"
            res_text = model.generate_s2s(wavs=feat, wavs_len=feat_lens, prompt=prompt)[0]
            token_list_str = res_text.split('|')[-1]
            token_list = json.loads(token_list_str)
            token_list = [int(i) for i in token_list]
            token_list2wav2(token_list, output_wav_path)


gpu_id = 7
def init_model_my():
    logging.basicConfig(level=logging.DEBUG,
                        format='%(asctime)s %(levelname)s %(message)s')
    checkpoint_path = "/mnt/sfs/asr/code/osum_xlgeng/examples/wenetspeech/whisper/exp/epoch_27_LLMbase_cosyvoice1_10Wtts_2Khqtts_3Ks2s_5Ws2t_streaming_s2s/step_119999.pt"
    config_path = "../conf/config_llm_huawei_base-version_cosyvoice1-token.yaml"
    args = GxlNode({
        "checkpoint": checkpoint_path,
    })
    configs = utils_file.load_dict_from_yaml(config_path)
    model, configs = init_model(args, configs)
    if is_npu:
        device = torch.device(f'npu:{gpu_id}')
    else:
        device =torch.device(f'cuda:{gpu_id}')
    model = model.to(device)
    tokenizer = init_tokenizer(configs)
    print(model)
    return model, tokenizer, device


model, tokenizer, device = init_model_my()
model.eval()
model = model.to(model.llama_model.dtype)


def do_resample(input_wav_path, output_wav_path):
    """"""
    waveform, sample_rate = torchaudio.load(input_wav_path)
    # 检查音频的维度
    num_channels = waveform.shape[0]
    # 如果音频是多通道的，则进行通道平均
    if num_channels > 1:
        waveform = torch.mean(waveform, dim=0, keepdim=True)
    waveform = torchaudio.transforms.Resample(
        orig_freq=sample_rate, new_freq=16000)(waveform)
    utils_file.makedir_for_file(output_wav_path)
    torchaudio.save(output_wav_path, waveform, 16000)

def get_feat_from_wav_path(input_wav_path, device=None):
    """
    获取音频的特征
    Args:
        input_wav_path: str

    Returns:
        feat: tensor, shape=(1, T, 80)
        feat_lens: tensor, shape=(1,)
    """
    timestamp_ms = int(time.time() * 1000)
    now_file_tmp_path_resample = f'~/.cache/.temp/{timestamp_ms}_resample.wav'
    do_resample(input_wav_path, now_file_tmp_path_resample)
    input_wav_path = now_file_tmp_path_resample
    waveform, sample_rate = torchaudio.load(input_wav_path)
    waveform = waveform.squeeze(0)  # (channel=1, sample) -> (sample,)
    window = torch.hann_window(400)
    stft = torch.stft(waveform,
                      400,
                      160,
                      window=window,
                      return_complex=True)
    magnitudes = stft[..., :-1].abs() ** 2

    filters = torch.from_numpy(
        librosa.filters.mel(sr=sample_rate,
                            n_fft=400,
                            n_mels=80))
    mel_spec = filters @ magnitudes

    # NOTE(xcsong): https://github.com/openai/whisper/discussions/269
    log_spec = torch.clamp(mel_spec, min=1e-10).log10()
    log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
    log_spec = (log_spec + 4.0) / 4.0
    feat = log_spec.transpose(0, 1).to(torch.bfloat16) # @wsy:转换数据类型为BFLOAT16
    feat_lens = torch.tensor([feat.shape[0]], dtype=torch.int64).to(device)
    feat = feat.unsqueeze(0).to(device)
    return feat, feat_lens


def infer_tts(input_text):
    prompt = "恳请将如下文本转换为其对应的语音token，力求生成最为流畅、自然的语音。"
    res_text = model.generate_tts(device=device, text=input_text, prompt=prompt)[0]
    return res_text
from cosyvoice_util import token_list2wav2
def infer_s2s(input_wav_path, output_wav_path):
    prompt = "先根据语音输入，直接以文字形式进行回答或对话，接着再生成语音token。"
    feat, feat_lens = get_feat_from_wav_path(input_wav_path,device)
    res_text = model.generate_s2s(wavs=feat, wavs_len=feat_lens, prompt=prompt)[0]
    token_list_str = res_text.split('|')[-1]
    token_list = json.loads(token_list_str)
    token_list = [int(i) for i in token_list]
    token_list2wav2(token_list, output_wav_path)
    return res_text

def infer_s2s_streaming(input_wav_path):
    prompt = "实现语音与文字的即时交流，直接进行文字和语音token的交错回复，每回复6个文字,就转化为18个语音token。"
    feat, feat_lens = get_feat_from_wav_path(input_wav_path)
    res_text = model.generate_s2s_streaming(wavs=feat, wavs_len=feat_lens, prompt=prompt)[0]
    return res_text
if __name__=="__main__":
    """"""
    infer_s2s("./input_data/question1.wav","./tmp.wav")
    # main()